{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Processed Data into Vector Database\n",
    "\n",
    "This notebook loads output from data prep kit into Milvus\n",
    "\n",
    "**Step-5 in this workflow**\n",
    "\n",
    "![](media/rag-overview-2.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-1: Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from my_config import MY_CONFIG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-2: Load Parquet Data\n",
    "\n",
    "Load all  `.parquet` files in the given dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data from :  output/output_final\n",
      "Number of parquet files to read :  2\n",
      "\n",
      "Read file: 'output/output_final/granite.parquet'.  number of rows = 33\n",
      "Read file: 'output/output_final/attention.parquet'.  number of rows = 28\n",
      "\n",
      "Total number of rows = 61\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import glob\n",
    "\n",
    "print ('Loading data from : ', MY_CONFIG.OUTPUT_FOLDER_FINAL)\n",
    "\n",
    "# Get a list of all Parquet files in the directory\n",
    "parquet_files = glob.glob(f'{MY_CONFIG.OUTPUT_FOLDER_FINAL}/*.parquet')\n",
    "print (\"Number of parquet files to read : \", len(parquet_files))\n",
    "print ()\n",
    "\n",
    "# Create an empty list to store the DataFrames\n",
    "dfs = []\n",
    "\n",
    "# Loop through each Parquet file and read it into a DataFrame\n",
    "for file in parquet_files:\n",
    "    df = pd.read_parquet(file)\n",
    "    print (f\"Read file: '{file}'.  number of rows = {df.shape[0]}\")\n",
    "    dfs.append(df)\n",
    "\n",
    "# Concatenate all DataFrames into a single DataFrame\n",
    "data_df = pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "print (f\"\\nTotal number of rows = {data_df.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding length:  384\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 61 entries, 0 to 60\n",
      "Data columns (total 15 columns):\n",
      " #   Column                 Non-Null Count  Dtype  \n",
      "---  ------                 --------------  -----  \n",
      " 0   filename               61 non-null     object \n",
      " 1   num_pages              61 non-null     int64  \n",
      " 2   num_tables             61 non-null     int64  \n",
      " 3   num_doc_elements       61 non-null     int64  \n",
      " 4   document_hash          61 non-null     object \n",
      " 5   ext                    61 non-null     object \n",
      " 6   hash                   61 non-null     object \n",
      " 7   size                   61 non-null     int64  \n",
      " 8   date_acquired          61 non-null     object \n",
      " 9   document_convert_time  61 non-null     float64\n",
      " 10  source_filename        61 non-null     object \n",
      " 11  source_document_id     61 non-null     object \n",
      " 12  text                   61 non-null     object \n",
      " 13  document_id            61 non-null     object \n",
      " 14  vector                 61 non-null     object \n",
      "dtypes: float64(1), int64(4), object(10)\n",
      "memory usage: 7.3+ KB\n",
      "None\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>filename</th>\n",
       "      <th>num_pages</th>\n",
       "      <th>num_tables</th>\n",
       "      <th>num_doc_elements</th>\n",
       "      <th>document_hash</th>\n",
       "      <th>ext</th>\n",
       "      <th>hash</th>\n",
       "      <th>size</th>\n",
       "      <th>date_acquired</th>\n",
       "      <th>document_convert_time</th>\n",
       "      <th>source_filename</th>\n",
       "      <th>source_document_id</th>\n",
       "      <th>text</th>\n",
       "      <th>document_id</th>\n",
       "      <th>vector</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>485</td>\n",
       "      <td>3127757990743433032</td>\n",
       "      <td>pdf</td>\n",
       "      <td>58342470e7d666dca0be87a15fb0552f949a5632606fe1...</td>\n",
       "      <td>121131</td>\n",
       "      <td>2025-10-03T09:32:24.162936</td>\n",
       "      <td>59.185795</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>adae9d79-6769-4eb7-96c0-b4021394dc25</td>\n",
       "      <td>## Granite Code Models: A Family of Open Found...</td>\n",
       "      <td>4ba39540df65ca93d9dc3026e7ffcea2f949ce9815229c...</td>\n",
       "      <td>[0.0064111887, -0.030633124, 0.005321988, 0.04...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>485</td>\n",
       "      <td>3127757990743433032</td>\n",
       "      <td>pdf</td>\n",
       "      <td>58342470e7d666dca0be87a15fb0552f949a5632606fe1...</td>\n",
       "      <td>121131</td>\n",
       "      <td>2025-10-03T09:32:24.162936</td>\n",
       "      <td>59.185795</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>adae9d79-6769-4eb7-96c0-b4021394dc25</td>\n",
       "      <td>## Abstract\\n\\nLarge Language Models (LLMs) tr...</td>\n",
       "      <td>fa843bc8b05ba80dd1ee780200b1614af27a2818759fa4...</td>\n",
       "      <td>[0.026825873, -0.026808968, 0.11900199, 0.0277...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>28</td>\n",
       "      <td>17</td>\n",
       "      <td>485</td>\n",
       "      <td>3127757990743433032</td>\n",
       "      <td>pdf</td>\n",
       "      <td>58342470e7d666dca0be87a15fb0552f949a5632606fe1...</td>\n",
       "      <td>121131</td>\n",
       "      <td>2025-10-03T09:32:24.162936</td>\n",
       "      <td>59.185795</td>\n",
       "      <td>granite.pdf</td>\n",
       "      <td>adae9d79-6769-4eb7-96c0-b4021394dc25</td>\n",
       "      <td>## 1 Introduction\\n\\nOver the last several dec...</td>\n",
       "      <td>d8469bf02cb5f03e01f372eb1c1265acf439855003d1a5...</td>\n",
       "      <td>[-0.029659774, 0.0077403677, 0.08020694, 0.006...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      filename  num_pages  num_tables  num_doc_elements        document_hash  \\\n",
       "0  granite.pdf         28          17               485  3127757990743433032   \n",
       "1  granite.pdf         28          17               485  3127757990743433032   \n",
       "2  granite.pdf         28          17               485  3127757990743433032   \n",
       "\n",
       "   ext                                               hash    size  \\\n",
       "0  pdf  58342470e7d666dca0be87a15fb0552f949a5632606fe1...  121131   \n",
       "1  pdf  58342470e7d666dca0be87a15fb0552f949a5632606fe1...  121131   \n",
       "2  pdf  58342470e7d666dca0be87a15fb0552f949a5632606fe1...  121131   \n",
       "\n",
       "                date_acquired  document_convert_time source_filename  \\\n",
       "0  2025-10-03T09:32:24.162936              59.185795     granite.pdf   \n",
       "1  2025-10-03T09:32:24.162936              59.185795     granite.pdf   \n",
       "2  2025-10-03T09:32:24.162936              59.185795     granite.pdf   \n",
       "\n",
       "                     source_document_id  \\\n",
       "0  adae9d79-6769-4eb7-96c0-b4021394dc25   \n",
       "1  adae9d79-6769-4eb7-96c0-b4021394dc25   \n",
       "2  adae9d79-6769-4eb7-96c0-b4021394dc25   \n",
       "\n",
       "                                                text  \\\n",
       "0  ## Granite Code Models: A Family of Open Found...   \n",
       "1  ## Abstract\\n\\nLarge Language Models (LLMs) tr...   \n",
       "2  ## 1 Introduction\\n\\nOver the last several dec...   \n",
       "\n",
       "                                         document_id  \\\n",
       "0  4ba39540df65ca93d9dc3026e7ffcea2f949ce9815229c...   \n",
       "1  fa843bc8b05ba80dd1ee780200b1614af27a2818759fa4...   \n",
       "2  d8469bf02cb5f03e01f372eb1c1265acf439855003d1a5...   \n",
       "\n",
       "                                              vector  \n",
       "0  [0.0064111887, -0.030633124, 0.005321988, 0.04...  \n",
       "1  [0.026825873, -0.026808968, 0.11900199, 0.0277...  \n",
       "2  [-0.029659774, 0.0077403677, 0.08020694, 0.006...  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "## Shape the data\n",
    "\n",
    "MY_CONFIG.EMBEDDING_LENGTH =  len(data_df.iloc[0]['embeddings'])\n",
    "print ('embedding length: ', MY_CONFIG.EMBEDDING_LENGTH)\n",
    "\n",
    "# rename 'embeddings' columns as 'vector' to match default schema\n",
    "# if 'vector' not in data_df.columns and 'embeddings' in data_df.columns:\n",
    "#     data_df = data_df.rename( columns= {'embeddings' : 'vector'})\n",
    "# if 'text' not in data_df.columns and 'contents' in data_df.columns:\n",
    "#     data_df = data_df.rename( columns= {'contents' : 'text'})\n",
    "\n",
    "data_df = data_df.rename( columns= {'embeddings' : 'vector', 'contents' : 'text'})\n",
    "\n",
    "print (data_df.info())\n",
    "data_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-3: Connect to Vector Database\n",
    "\n",
    "Milvus can be embedded and easy to use.\n",
    "\n",
    "<span style=\"color:blue;\">Note: If you encounter an error about unable to load database, try this: </span>\n",
    "\n",
    "- <span style=\"color:blue;\">In **vscode** : **restart the kernel** of previous notebook. This will release the db.lock </span>\n",
    "- <span style=\"color:blue;\">In **Jupyter**: Do `File --> Close and Shutdown Notebook` of previous notebook. This will release the db.lock</span>\n",
    "- <span style=\"color:blue;\">Re-run this cell again</span>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/shalisha.witherspoonibm.com/Documents/DPK_RAG_UPDATE/data-prep-kit-outer/examples/rag-pdf-1/venv/lib/python3.12/site-packages/milvus_lite/__init__.py:15: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  from pkg_resources import DistributionNotFound, get_distribution\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Connected to Milvus instance: ./rag_1_dpk.db\n"
     ]
    }
   ],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(MY_CONFIG.DB_URI)\n",
    "\n",
    "print (\"✅ Connected to Milvus instance:\", MY_CONFIG.DB_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step-4: Create A Collection\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Created collection : dpk_papers\n"
     ]
    }
   ],
   "source": [
    "# if we already have a collection, clear it first\n",
    "if milvus_client.has_collection(collection_name=MY_CONFIG.COLLECTION_NAME):\n",
    "    milvus_client.drop_collection(collection_name=MY_CONFIG.COLLECTION_NAME)\n",
    "    print ('✅ Cleared collection :', MY_CONFIG.COLLECTION_NAME)\n",
    "\n",
    "\n",
    "milvus_client.create_collection(\n",
    "    collection_name=MY_CONFIG.COLLECTION_NAME,\n",
    "    dimension=MY_CONFIG.EMBEDDING_LENGTH,\n",
    "    metric_type=\"IP\",  # Inner product distance\n",
    "    consistency_level=\"Strong\",  # Strong consistency level\n",
    "    auto_id=True\n",
    ")\n",
    "print (\"✅ Created collection :\", MY_CONFIG.COLLECTION_NAME)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-5: Insert Data into Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inserted # rows 61\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'row_count': 61}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = milvus_client.insert(collection_name=MY_CONFIG.COLLECTION_NAME, data=data_df.to_dict('records'))\n",
    "\n",
    "print('inserted # rows', res['insert_count'])\n",
    "\n",
    "milvus_client.get_collection_stats(MY_CONFIG.COLLECTION_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step-6: Close DB Connection\n",
    "\n",
    "Close the connection so the lock files are relinquished and other notebooks can access the db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ SUCCESS\n"
     ]
    }
   ],
   "source": [
    "milvus_client.close()\n",
    "\n",
    "print (\"✅ SUCCESS\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test your data by doing a Vector Search\n",
    "\n",
    "See notebook [vector_search.ipynb](vector_search.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
